#encoding: utf8

from scapy.all import *
from netaddr import IPNetwork,IPAddress
import redis
from search_ip import ip2host
 

rds = redis.Redis(host='127.0.0.1', port=6379,db=0)
school_ip = [str(ipaddr) for ipaddr in IPNetwork("192.168.1.0/24")]
# school_ip = [str(ipaddr) for ipaddr in IPNetwork("125.216.241.0/24")]
school_ip = set(school_ip)
visit_dict = {}

def increase_visit_cnt(src2dst):
    cnt_str = rds.get(src2dst)
    if cnt_str:
        cnt = int(cnt_str)
        cnt += 1
        rds.set(src2dst, cnt)
    else:
        cnt = 1
        rds.set(src2dst, cnt)

def store_ipsrc_to_redis(pkt):

    if not pkt.haslayer("IP"):
        return ""

    ip_src = pkt[IP].src
    ip_dst = pkt[IP].dst

    if ip_src in school_ip:
        rds.lpush(ip_src, "%s | %s" %(ip_dst, pkt.summary()))
        src2dst = "%s->%s" %(ip_src, ip_dst)
        increase_visit_cnt(src2dst)

        uri = ip2host(ip_dst)
        src2dst = "%s->%s" %(ip_src, uri)
        return src2dst

    elif ip_dst in school_ip:
        rds.lpush(ip_dst, "%s | %s" %(ip_src, pkt.summary()))
    else:
        rds.lpush("other", "%s" %( pkt.summary()))

def store_dns_record(pkt):
    if pkt.haslayer("DNS"):

        # DNS ans packet
        if pkt[DNS].qr == 1:
            
            cnt = int(pkt[DNS].ancount) 
            if cnt >= 1:
                uri = pkt[DNS].an[0].rrname
                ip_list = []
                for i in range(cnt):
                    # A record
                    if pkt[DNS].an[i].type == 1:
                        ip_list.append(pkt[DNS].an[i].rdata)
                
                        key = "DNS_%s" % pkt[DNS].an[i].rdata
                        val = uri
                        rds.set(key, val)
                return "uri: %s, ip: %s" %(uri, str(ip_list))

            else: 
                pass


def packet_handler(pkt): 
    if pkt.haslayer("DNS"):
        return store_dns_record(pkt)

    elif pkt.haslayer("TCP"):
        return store_ipsrc_to_redis(pkt)

    


# all_packets = rdpcap("image_http_request.pcap")
# idx = 0
# for pkt in all_packets:
#     idx += 1
#     packet_handler(pkt)  
# pkts =  sniff(prn = lambda x: packet_handler(x), store=0, iface="eth0")
# pkts = sniff(prn=lambda x:x.sprintf("{IP:%IP.src% -> %IP.dst%\n}{Raw:%Raw.load%\n}"))
pkts = sniff(prn=lambda x: packet_handler(x), store=0)


print "never come here"
